"""
Background removal functionality for images.
"""

import io
import base64
from typing import Optional
from PIL import Image
import logging

logger = logging.getLogger(__name__)

try:
    from rembg import remove, new_session
    REMBG_AVAILABLE = True
    logger.info("rembg library loaded successfully")
except ImportError as e:
    REMBG_AVAILABLE = False
    # Store the import error for debugging
    _REMBG_IMPORT_ERROR = str(e)
    logger.warning(f"rembg library not available: {e}")
except Exception as e:
    # Catch other potential errors like NumPy version conflicts
    REMBG_AVAILABLE = False
    _REMBG_IMPORT_ERROR = f"Error loading rembg: {e}"
    logger.warning(f"Error loading rembg: {e}")


class BackgroundRemover:
    """Handle background removal from images."""
    
    def __init__(self):
        """Initialize the background remover."""
        if not REMBG_AVAILABLE:
            error_msg = f"rembg library not available. Background removal will not work."
            if '_REMBG_IMPORT_ERROR' in globals():
                error_msg += f" Import error: {_REMBG_IMPORT_ERROR}"
            logger.warning(error_msg)
        
        self._sessions = {}
    
    def _get_session(self, model_name: str):
        """Get or create a rembg session for the specified model.
        
        Args:
            model_name: Name of the background removal model
            
        Returns:
            rembg session
        """
        if model_name not in self._sessions:
            self._sessions[model_name] = new_session(model_name)
        return self._sessions[model_name]
    
    async def remove_background(self, image_b64: str, model_name: str = "u2net") -> str:
        """Remove background from image and return as base64 PNG with transparency.
        
        Args:
            image_b64: Base64 encoded input image
            model_name: Background removal model to use
            
        Returns:
            Base64 encoded PNG image with transparent background
        """
        if not REMBG_AVAILABLE:
            raise Exception("rembg library is required for background removal. Please install it with: pip install rembg")
        
        try:
            logger.info(f"Removing background using model: {model_name}")
            
            # Decode base64 image
            image_bytes = base64.b64decode(image_b64)
            
            # Load image
            input_image = Image.open(io.BytesIO(image_bytes))
            
            # Get rembg session
            session = self._get_session(model_name)
            
            # Remove background
            output_image = remove(input_image, session=session)
            
            # Ensure the output is RGBA (with alpha channel)
            if output_image.mode != 'RGBA':
                output_image = output_image.convert('RGBA')
            
            # Save to bytes buffer as PNG (which supports transparency)
            output_buffer = io.BytesIO()
            output_image.save(output_buffer, format='PNG')
            output_bytes = output_buffer.getvalue()
            
            # Encode to base64
            output_b64 = base64.b64encode(output_bytes).decode('utf-8')
            
            logger.info("Background removed successfully")
            return output_b64
            
        except Exception as e:
            logger.error(f"Error removing background: {str(e)}")
            raise Exception(f"Failed to remove background: {str(e)}")
